/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.beam.sdk.schemas.utils;

import static org.apache.beam.sdk.util.ByteBuddyUtils.getClassLoadingStrategy;

import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;
import java.util.stream.Collectors;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.asm.AsmVisitorWrapper;
import net.bytebuddy.description.method.MethodDescription.ForLoadedMethod;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.dynamic.scaffold.InstrumentedType;
import net.bytebuddy.implementation.FixedValue;
import net.bytebuddy.implementation.Implementation;
import net.bytebuddy.implementation.bytecode.ByteCodeAppender;
import net.bytebuddy.implementation.bytecode.ByteCodeAppender.Size;
import net.bytebuddy.implementation.bytecode.Removal;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.assign.TypeCasting;
import net.bytebuddy.implementation.bytecode.member.MethodInvocation;
import net.bytebuddy.implementation.bytecode.member.MethodReturn;
import net.bytebuddy.implementation.bytecode.member.MethodVariableAccess;
import net.bytebuddy.jar.asm.ClassWriter;
import net.bytebuddy.matcher.ElementMatchers;
import org.apache.beam.sdk.schemas.FieldValueGetter;
import org.apache.beam.sdk.schemas.FieldValueHaver;
import org.apache.beam.sdk.schemas.FieldValueSetter;
import org.apache.beam.sdk.schemas.FieldValueTypeInformation;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.SchemaUserTypeCreator;
import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.ConstructorCreateInstruction;
import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.InjectPackageStrategy;
import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.StaticFactoryMethodInstruction;
import org.apache.beam.sdk.schemas.utils.ByteBuddyUtils.TypeConversionsFactory;
import org.apache.beam.sdk.schemas.utils.ReflectUtils.TypeDescriptorWithSchema;
import org.apache.beam.sdk.util.Preconditions;
import org.apache.beam.sdk.util.common.ReflectHelpers;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
import org.checkerframework.checker.nullness.qual.NonNull;
import org.checkerframework.checker.nullness.qual.Nullable;

/** A set of utilities to generate getter and setter classes for JavaBean objects. */
@SuppressWarnings({"rawtypes"})
public class JavaBeanUtils {

  private static final String X_WITH_NULL_METHOD_ERROR_FMT =
      "a %s FieldValueTypeInformation object has a null method field";
  public static final String GETTER_WITH_NULL_METHOD_ERROR =
      String.format(X_WITH_NULL_METHOD_ERROR_FMT, "getter");
  public static final String SETTER_WITH_NULL_METHOD_ERROR =
      String.format(X_WITH_NULL_METHOD_ERROR_FMT, "setter");

  /** Create a {@link Schema} for a Java Bean class. */
  public static Schema schemaFromJavaBeanClass(
      TypeDescriptor<?> typeDescriptor, FieldValueTypeSupplier fieldValueTypeSupplier) {
    return StaticSchemaInference.schemaFromClass(typeDescriptor, fieldValueTypeSupplier);
  }

  private static final String CONSTRUCTOR_HELP_STRING =
      "In order to infer a Schema from a Java Bean, it must have a constructor annotated with"
          + " @SchemaCreate, or it must have a compatible setter for every getter used as a Schema"
          + " field.";

  // Make sure that there are matching setters and getters.
  public static void validateJavaBean(
      List<FieldValueTypeInformation> getters,
      List<FieldValueTypeInformation> setters,
      Schema schema) {

    Map<String, FieldValueTypeInformation> setterMap = new HashMap<>();
    int bound = schema.getFieldCount();
    for (int i = 0; i < bound; i++) {
      Integer integer = i;
      if (setterMap.put(schema.getField(integer).getName(), setters.get(integer)) != null) {
        throw new IllegalStateException("Duplicate key");
      }
    }

    for (FieldValueTypeInformation type : getters) {
      FieldValueTypeInformation setterType = setterMap.get(type.getName());
      Method m =
          Preconditions.checkArgumentNotNull(type.getMethod(), GETTER_WITH_NULL_METHOD_ERROR);
      if (setterType == null) {
        throw new RuntimeException(
            String.format(
                "Java Bean '%s' contains a getter for field '%s', but does not contain a matching"
                    + " setter. %s",
                m.getDeclaringClass(), type.getName(), CONSTRUCTOR_HELP_STRING));
      }
      if (!type.getType().equals(setterType.getType())) {
        throw new RuntimeException(
            String.format(
                "Java Bean '%s' contains a setter for field '%s' that has a mismatching type. %s",
                m.getDeclaringClass(), type.getName(), CONSTRUCTOR_HELP_STRING));
      }
      if (!type.isNullable() == setterType.isNullable()) {
        throw new RuntimeException(
            String.format(
                "Java Bean '%s' contains a setter for field '%s' that has a mismatching nullable"
                    + " attribute. %s",
                m.getDeclaringClass(), type.getName(), CONSTRUCTOR_HELP_STRING));
      }
    }
  }

  // Static ByteBuddy instance used by all helpers.
  private static final ByteBuddy BYTE_BUDDY = new ByteBuddy();

  private static final Map<TypeDescriptorWithSchema<?>, List<FieldValueTypeInformation>>
      CACHED_FIELD_TYPES = Maps.newConcurrentMap();

  public static List<FieldValueTypeInformation> getFieldTypes(
      TypeDescriptor<?> typeDescriptor,
      Schema schema,
      FieldValueTypeSupplier fieldValueTypeSupplier) {
    return CACHED_FIELD_TYPES.computeIfAbsent(
        TypeDescriptorWithSchema.create(typeDescriptor, schema),
        c -> fieldValueTypeSupplier.get(typeDescriptor, schema));
  }

  // The list of getters for a class is cached, so we only create the classes the first time
  // getSetters is called.
  private static final Map<TypeDescriptorWithSchema<?>, List<FieldValueGetter<@NonNull ?, ?>>>
      CACHED_GETTERS = Maps.newConcurrentMap();

  /**
   * Return the list of {@link FieldValueGetter}s for a Java Bean class
   *
   * <p>The returned list is ordered by the order of fields in the schema.
   */
  public static <T> List<FieldValueGetter<@NonNull T, Object>> getGetters(
      TypeDescriptor<T> typeDescriptor,
      Schema schema,
      FieldValueTypeSupplier fieldValueTypeSupplier,
      TypeConversionsFactory typeConversionsFactory) {
    return (List)
        CACHED_GETTERS.computeIfAbsent(
            TypeDescriptorWithSchema.create(typeDescriptor, schema),
            c -> {
              List<FieldValueTypeInformation> types =
                  fieldValueTypeSupplier.get(typeDescriptor, schema);
              return types.stream()
                  .map(t -> JavaBeanUtils.createGetter(t, typeConversionsFactory))
                  .collect(Collectors.toList());
            });
  }

  public static <ObjectT extends @NonNull Object, ValueT>
      FieldValueGetter<ObjectT, ValueT> createGetter(
          FieldValueTypeInformation typeInformation,
          TypeConversionsFactory typeConversionsFactory) {
    final Method m =
        Preconditions.checkArgumentNotNull(
            typeInformation.getMethod(), GETTER_WITH_NULL_METHOD_ERROR);
    DynamicType.Builder<FieldValueGetter<ObjectT, ValueT>> builder =
        ByteBuddyUtils.subclassGetterInterface(
            BYTE_BUDDY,
            m.getDeclaringClass(),
            typeConversionsFactory.createTypeConversion(false).convert(typeInformation.getType()));
    builder = implementGetterMethods(builder, typeInformation, typeConversionsFactory);
    try {
      return builder
          .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES))
          .make()
          .load(
              ReflectHelpers.findClassLoader(m.getDeclaringClass().getClassLoader()),
              getClassLoadingStrategy(m.getDeclaringClass()))
          .getLoaded()
          .getDeclaredConstructor()
          .newInstance();
    } catch (InstantiationException
        | IllegalAccessException
        | NoSuchMethodException
        | InvocationTargetException e) {
      throw new RuntimeException(
          "Unable to generate a getter for getter '" + typeInformation.getMethod() + "'");
    }
  }

  private static <ObjectT extends @NonNull Object, ValueT>
      DynamicType.Builder<FieldValueGetter<ObjectT, ValueT>> implementGetterMethods(
          DynamicType.Builder<FieldValueGetter<ObjectT, ValueT>> builder,
          FieldValueTypeInformation typeInformation,
          TypeConversionsFactory typeConversionsFactory) {
    return builder
        .method(ElementMatchers.named("name"))
        .intercept(FixedValue.reference(typeInformation.getName()))
        .method(ElementMatchers.named("get"))
        .intercept(new InvokeGetterInstruction(typeInformation, typeConversionsFactory));
  }

  // The list of setters for a class is cached, so we only create the classes the first time
  // getSetters is called.
  private static final Map<TypeDescriptorWithSchema<?>, List<FieldValueSetter>> CACHED_SETTERS =
      Maps.newConcurrentMap();

  /**
   * Return the list of {@link FieldValueSetter}s for a Java Bean class
   *
   * <p>The returned list is ordered by the order of fields in the schema.
   */
  public static List<FieldValueSetter> getSetters(
      TypeDescriptor<?> typeDescriptor,
      Schema schema,
      FieldValueTypeSupplier fieldValueTypeSupplier,
      TypeConversionsFactory typeConversionsFactory) {
    return CACHED_SETTERS.computeIfAbsent(
        TypeDescriptorWithSchema.create(typeDescriptor, schema),
        c -> {
          List<FieldValueTypeInformation> types =
              fieldValueTypeSupplier.get(typeDescriptor, schema);
          return types.stream()
              .map(t -> createSetter(t, typeConversionsFactory))
              .collect(Collectors.toList());
        });
  }

  public static <ObjectT, ValueT> FieldValueSetter<ObjectT, ValueT> createSetter(
      FieldValueTypeInformation typeInformation, TypeConversionsFactory typeConversionsFactory) {
    final Method m =
        Preconditions.checkArgumentNotNull(
            typeInformation.getMethod(), SETTER_WITH_NULL_METHOD_ERROR);
    DynamicType.Builder<FieldValueSetter<ObjectT, ValueT>> builder =
        ByteBuddyUtils.subclassSetterInterface(
            BYTE_BUDDY,
            m.getDeclaringClass(),
            typeConversionsFactory.createTypeConversion(false).convert(typeInformation.getType()));
    builder = implementSetterMethods(builder, typeInformation, typeConversionsFactory);
    try {
      return builder
          .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES))
          .make()
          .load(
              ReflectHelpers.findClassLoader(m.getDeclaringClass().getClassLoader()),
              getClassLoadingStrategy(m.getDeclaringClass()))
          .getLoaded()
          .getDeclaredConstructor()
          .newInstance();
    } catch (InstantiationException
        | IllegalAccessException
        | NoSuchMethodException
        | InvocationTargetException e) {
      throw new RuntimeException(
          "Unable to generate a setter for setter '" + typeInformation.getMethod() + "'");
    }
  }

  private static <ObjectT, ValueT>
      DynamicType.Builder<FieldValueSetter<ObjectT, ValueT>> implementSetterMethods(
          DynamicType.Builder<FieldValueSetter<ObjectT, ValueT>> builder,
          FieldValueTypeInformation fieldValueTypeInformation,
          TypeConversionsFactory typeConversionsFactory) {
    return builder
        .method(ElementMatchers.named("name"))
        .intercept(FixedValue.reference(fieldValueTypeInformation.getName()))
        .method(ElementMatchers.named("set"))
        .intercept(new InvokeSetterInstruction(fieldValueTypeInformation, typeConversionsFactory));
  }

  public static <ObjectT> FieldValueHaver<ObjectT> createHaver(
      Class<ObjectT> clazz, Method hasMethod) {
    DynamicType.Builder<FieldValueHaver<ObjectT>> builder =
        ByteBuddyUtils.subclassHaverInterface(BYTE_BUDDY, clazz);
    builder = implementHaverMethods(builder, hasMethod);
    try {
      return builder
          .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES))
          .make()
          .load(
              ReflectHelpers.findClassLoader(clazz.getClassLoader()),
              getClassLoadingStrategy(clazz))
          .getLoaded()
          .getDeclaredConstructor()
          .newInstance();
    } catch (InstantiationException
        | IllegalAccessException
        | InvocationTargetException
        | NoSuchMethodException e) {
      throw new RuntimeException("Unable to generate a have for hasMethod '" + hasMethod + "'", e);
    }
  }

  private static <ObjectT> DynamicType.Builder<FieldValueHaver<ObjectT>> implementHaverMethods(
      DynamicType.Builder<FieldValueHaver<ObjectT>> builder, Method hasMethod) {
    return builder
        .method(ElementMatchers.named("name"))
        .intercept(FixedValue.reference(hasMethod.getName()))
        .method(ElementMatchers.named("has"))
        .intercept(new InvokeHaverInstruction(hasMethod));
  }

  // The list of constructors for a class is cached, so we only create the classes the first time
  // getConstructor is called.
  public static final Map<TypeDescriptorWithSchema<?>, SchemaUserTypeCreator> CACHED_CREATORS =
      Maps.newConcurrentMap();

  public static SchemaUserTypeCreator getConstructorCreator(
      TypeDescriptor<?> typeDescriptor,
      Constructor constructor,
      Schema schema,
      FieldValueTypeSupplier fieldValueTypeSupplier,
      TypeConversionsFactory typeConversionsFactory) {
    return CACHED_CREATORS.computeIfAbsent(
        TypeDescriptorWithSchema.create(typeDescriptor, schema),
        c -> {
          List<FieldValueTypeInformation> types =
              fieldValueTypeSupplier.get(typeDescriptor, schema);
          return createConstructorCreator(
              typeDescriptor.getRawType(), constructor, schema, types, typeConversionsFactory);
        });
  }

  public static <T> SchemaUserTypeCreator createConstructorCreator(
      Class<T> clazz,
      Constructor<T> constructor,
      Schema schema,
      List<FieldValueTypeInformation> types,
      TypeConversionsFactory typeConversionsFactory) {
    try {
      DynamicType.Builder<SchemaUserTypeCreator> builder =
          BYTE_BUDDY
              .with(new InjectPackageStrategy(clazz))
              .subclass(SchemaUserTypeCreator.class)
              .method(ElementMatchers.named("create"))
              .intercept(
                  new ConstructorCreateInstruction(
                      types, clazz, constructor, typeConversionsFactory));
      return builder
          .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES))
          .make()
          .load(
              ReflectHelpers.findClassLoader(clazz.getClassLoader()),
              getClassLoadingStrategy(clazz))
          .getLoaded()
          .getDeclaredConstructor()
          .newInstance();
    } catch (InstantiationException
        | IllegalAccessException
        | NoSuchMethodException
        | InvocationTargetException e) {
      throw new RuntimeException(
          "Unable to generate a creator for class " + clazz + " with schema " + schema);
    }
  }

  public static SchemaUserTypeCreator getStaticCreator(
      TypeDescriptor<?> typeDescriptor,
      Method creator,
      Schema schema,
      FieldValueTypeSupplier fieldValueTypeSupplier,
      TypeConversionsFactory typeConversionsFactory) {
    return CACHED_CREATORS.computeIfAbsent(
        TypeDescriptorWithSchema.create(typeDescriptor, schema),
        c -> {
          List<FieldValueTypeInformation> types =
              fieldValueTypeSupplier.get(typeDescriptor, schema);
          return createStaticCreator(
              typeDescriptor.getRawType(), creator, schema, types, typeConversionsFactory);
        });
  }

  public static <T> SchemaUserTypeCreator createStaticCreator(
      Class<T> clazz,
      Method creator,
      Schema schema,
      List<FieldValueTypeInformation> types,
      TypeConversionsFactory typeConversionsFactory) {
    try {
      DynamicType.Builder<SchemaUserTypeCreator> builder =
          BYTE_BUDDY
              .with(new InjectPackageStrategy(clazz))
              .subclass(SchemaUserTypeCreator.class)
              .method(ElementMatchers.named("create"))
              .intercept(
                  new StaticFactoryMethodInstruction(
                      types, clazz, creator, typeConversionsFactory));

      return builder
          .visit(new AsmVisitorWrapper.ForDeclaredMethods().writerFlags(ClassWriter.COMPUTE_FRAMES))
          .make()
          .load(
              ReflectHelpers.findClassLoader(clazz.getClassLoader()),
              getClassLoadingStrategy(clazz))
          .getLoaded()
          .getDeclaredConstructor()
          .newInstance();
    } catch (InstantiationException
        | IllegalAccessException
        | NoSuchMethodException
        | InvocationTargetException e) {
      throw new RuntimeException(
          "Unable to generate a creator for " + clazz + " with schema " + schema);
    }
  }

  public static <T, K extends Comparable<? super K>> Comparator<T> comparingNullFirst(
      Function<? super T, ? extends @Nullable K> keyExtractor) {
    return Comparator.comparing(keyExtractor, Comparator.nullsFirst(Comparator.naturalOrder()));
  }

  // Implements a method to read a public getter out of an object.
  private static class InvokeGetterInstruction implements Implementation {
    private final FieldValueTypeInformation typeInformation;
    private final TypeConversionsFactory typeConversionsFactory;

    InvokeGetterInstruction(
        FieldValueTypeInformation typeInformation, TypeConversionsFactory typeConversionsFactory) {
      this.typeInformation = typeInformation;
      this.typeConversionsFactory = typeConversionsFactory;
    }

    @Override
    public InstrumentedType prepare(InstrumentedType instrumentedType) {
      return instrumentedType;
    }

    @Override
    public ByteCodeAppender appender(final Target implementationTarget) {
      return (methodVisitor, implementationContext, instrumentedMethod) -> {
        // this + method parameters.
        int numLocals = 1 + instrumentedMethod.getParameters().size();
        StackManipulation cast =
            typeInformation
                    .getRawType()
                    .isAssignableFrom(
                        Preconditions.checkStateNotNull(typeInformation.getMethod())
                            .getReturnType())
                ? StackManipulation.Trivial.INSTANCE
                : TypeCasting.to(TypeDescription.ForLoadedType.of(typeInformation.getRawType()));

        // StackManipulation that will read the value from the class field.
        StackManipulation readValue =
            new StackManipulation.Compound(
                // Method param is offset 1 (offset 0 is the this parameter).
                MethodVariableAccess.REFERENCE.loadFrom(1),
                // Invoke the getter
                MethodInvocation.invoke(
                    new ForLoadedMethod(
                        Preconditions.checkStateNotNull(
                            typeInformation.getMethod(), GETTER_WITH_NULL_METHOD_ERROR))),
                cast);

        StackManipulation stackManipulation =
            new StackManipulation.Compound(
                typeConversionsFactory
                    .createGetterConversions(readValue)
                    .convert(typeInformation.getType()),
                MethodReturn.REFERENCE);

        StackManipulation.Size size = stackManipulation.apply(methodVisitor, implementationContext);
        return new Size(size.getMaximalSize(), numLocals);
      };
    }
  }

  // Implements a method to write a public set out on an object.
  private static class InvokeSetterInstruction implements Implementation {
    // Setter method that will be invoked
    private FieldValueTypeInformation fieldValueTypeInformation;
    private final TypeConversionsFactory typeConversionsFactory;

    InvokeSetterInstruction(
        FieldValueTypeInformation fieldValueTypeInformation,
        TypeConversionsFactory typeConversionsFactory) {
      this.fieldValueTypeInformation = fieldValueTypeInformation;
      this.typeConversionsFactory = typeConversionsFactory;
    }

    @Override
    public InstrumentedType prepare(InstrumentedType instrumentedType) {
      return instrumentedType;
    }

    @Override
    public ByteCodeAppender appender(final Target implementationTarget) {
      return (methodVisitor, implementationContext, instrumentedMethod) -> {
        // this + method parameters.
        int numLocals = 1 + instrumentedMethod.getParameters().size();

        // The instruction to read the field.
        StackManipulation readField = MethodVariableAccess.REFERENCE.loadFrom(2);

        Method method =
            Preconditions.checkStateNotNull(
                fieldValueTypeInformation.getMethod(), SETTER_WITH_NULL_METHOD_ERROR);
        boolean setterMethodReturnsVoid = method.getReturnType().equals(Void.TYPE);
        // Read the object onto the stack.
        StackManipulation stackManipulation =
            new StackManipulation.Compound(
                // Object param is offset 1.
                MethodVariableAccess.REFERENCE.loadFrom(1),
                // Do any conversions necessary.
                typeConversionsFactory
                    .createSetterConversions(readField)
                    .convert(fieldValueTypeInformation.getType()),
                // Now update the field and return void.
                MethodInvocation.invoke(new ForLoadedMethod(method)));
        if (!setterMethodReturnsVoid) {
          // Discard return type;
          stackManipulation = new StackManipulation.Compound(stackManipulation, Removal.SINGLE);
        }
        stackManipulation = new StackManipulation.Compound(stackManipulation, MethodReturn.VOID);

        StackManipulation.Size size = stackManipulation.apply(methodVisitor, implementationContext);
        return new Size(size.getMaximalSize(), numLocals);
      };
    }
  }

  // Implements a method to check a presence on an object.
  private static class InvokeHaverInstruction implements Implementation {
    private final Method hasMethod;

    public InvokeHaverInstruction(Method hasMethod) {
      this.hasMethod = hasMethod;
    }

    @Override
    public ByteCodeAppender appender(Target implementationTarget) {
      return (methodVisitor, implementationContext, instrumentedMethod) -> {
        // this + method parameters.
        int numLocals = 1 + instrumentedMethod.getParameters().size();
        StackManipulation.Size size =
            new StackManipulation.Compound(
                    // Read the first argument
                    MethodVariableAccess.REFERENCE.loadFrom(1),
                    // Call hasMethod
                    MethodInvocation.invoke(new ForLoadedMethod(hasMethod)),
                    MethodReturn.INTEGER)
                .apply(methodVisitor, implementationContext);
        return new Size(size.getMaximalSize(), numLocals);
      };
    }

    @Override
    public InstrumentedType prepare(InstrumentedType instrumentedType) {
      return instrumentedType;
    }
  }
}
